#include <ipa_room_segmentation/room_segmentation_test.h>
#include <ipa_room_segmentation/meanshift2d.h>
#include <boost/algorithm/string.hpp>
#include <opencv2/imgproc/types_c.h>

const std::string package_path = "/home/ycz/ws/ipa_coverage_planning/src/ipa_room_segmentation/";
const std::string classifier_default_path = package_path + "common/files/classifier_models/";
const std::string classifier_path = "room_segmentation/classifier_models/";
std::string classifier_storage_path = "room_segmentation/classifier_models/";
std::map<int, std::string> officeImgNames{
        {0, "big"},
        {1, "office"},
        {2, "pass"},
        {3, "tea"},
        {4, "home"},
        {5, "origin_map"},
};

int main(int argc, char **argv)
{
    int room_segmentation_algorithm_ = 2;
    // displays the segmented map upon service call
    bool display_segmented_map_ = true;
    static bool DEBUG_DISPLAYS = true;
    std::cout << "\n------------------------------------\nRoom Segmentation Parameters:\n------------------------------------\n";
    std::cout << "room_segmentation/room_segmentation_algorithm = " << room_segmentation_algorithm_ << std::endl;
    std::cout << "room_segmentation/display_segmented_map_ = " << display_segmented_map_ << std::endl;
    //set the resolution and the limits for the actual goal and the Map origin
    const float map_resolution = 0.05;
    const float robot_radius = 0.05;
    std::string fileName = officeImgNames.at(4);
    std::string srcName = package_path + "img/" + fileName + ".png";
    cv::Point2d map_origin(0, 0);
    if (fileName == "big")map_origin = cv::Point2d(190, 120);
    if (fileName == "office")map_origin = cv::Point2d(830, 220);
    if (fileName == "pass")map_origin = cv::Point2d(670, 70);
    if (fileName == "tea")map_origin = cv::Point2d(365, 255);
    if (fileName == "home")map_origin = cv::Point2d(21, 79);
    if (fileName == "origin_map")map_origin = cv::Point2d(157, 17);
    // segment the given map
    std::cout << "\033[1;32msrcName: " << srcName << "\033[0m" << std::endl;
    cv::Mat segmented_map;
    cv::Mat original_img = cv::imread(srcName,cv::IMREAD_GRAYSCALE);
    imshow("original_img", original_img);
    cv::waitKey(1);
    // set morphological parameters
    if (room_segmentation_algorithm_ == 1)
    {
        room_upper_limit_morphological_ = 13.0;
        room_lower_limit_morphological_ = 0.3;
        std::cout << "You have chosen the morphological segmentation method."<< std::endl;
        std::cout << "room_segmentation/room_area_factor_upper_limit = "
                  << room_upper_limit_morphological_ << std::endl;
        std::cout << "room_segmentation/room_area_factor_lower_limit = "
                  << room_lower_limit_morphological_ << std::endl;
        MorphologicalSegmentation morphological_segmentation;
        morphological_segmentation.segmentMap(original_img, segmented_map, map_resolution,
                                              room_lower_limit_morphological_, room_upper_limit_morphological_);
    }
    // set distance parameters
    if (room_segmentation_algorithm_ == 2)
    {
        room_upper_limit_distance_ = 163.0;
        room_lower_limit_distance_ = 0.35;
        std::cout << "You have chosen the distance segmentation method."<< std::endl;
        std::cout << "room_segmentation/room_area_factor_upper_limit = " << room_upper_limit_distance_ << std::endl;
        std::cout << "room_segmentation/room_area_factor_lower_limit = " << room_lower_limit_distance_ << std::endl;
        DistanceSegmentation distance_segmentation;
        distance_segmentation.segmentMap(original_img, segmented_map, map_resolution, room_lower_limit_distance_,
                                         room_upper_limit_distance_);
    }
    // set voronoi parameters
    if (room_segmentation_algorithm_ == 3)
    {
        room_upper_limit_voronoi_ = 120.0;
        room_lower_limit_voronoi_ = 1.53;
        voronoi_neighborhood_index_ = 280;
        max_iterations_ = 150;
        min_critical_point_distance_factor_ = 1.6;
        max_area_for_merging_ = 12.5;
        std::cout << "You have chosen the voronoi segmentation method."<< std::endl;
        std::cout << "room_segmentation/room_area_factor_upper_limit = " << room_upper_limit_voronoi_ << std::endl;
        std::cout << "room_segmentation/room_area_factor_lower_limit = " << room_lower_limit_voronoi_ << std::endl;
        std::cout << "room_segmentation/voronoi_neighborhood_index = " << voronoi_neighborhood_index_ << std::endl;
        std::cout << "room_segmentation/max_iterations = " << max_iterations_ << std::endl;
        std::cout << "room_segmentation/min_critical_point_distance_factor = "
                  << min_critical_point_distance_factor_ << std::endl;
        std::cout << "room_segmentation/max_area_for_merging = " << max_area_for_merging_ << std::endl;
        VoronoiSegmentation voronoi_segmentation;
        voronoi_segmentation.segmentMap(original_img, segmented_map, map_resolution, room_lower_limit_voronoi_,
                                        room_upper_limit_voronoi_,
                                        voronoi_neighborhood_index_, max_iterations_,
                                        min_critical_point_distance_factor_, max_area_for_merging_,
                                        (display_segmented_map_ && DEBUG_DISPLAYS));
    }
    // set semantic parameters
    if (room_segmentation_algorithm_ == 4 || train_semantic_)
    {
        room_upper_limit_semantic_ = 23.0;
        room_lower_limit_semantic_ = 1.0;
        std::cout << "You have chosen the semantic segmentation method."<< std::endl;
        std::cout << "room_segmentation/room_area_factor_upper_limit = " << room_upper_limit_semantic_ << std::endl;
        std::cout << "room_segmentation/room_area_factor_lower_limit = " << room_lower_limit_semantic_ << std::endl;
        // train the algorithm if wanted
        if (train_semantic_)
        {
            AdaboostClassifier semantic_segmentation;
            std::cout << "You have chosen the voronoi random field segmentation method."<< std::endl;
            std::cout << "room_segmentation/semantic_training_maps_room_file_list_ = \n";
            std::vector <cv::Mat> room_training_maps;
            for (const auto & i : semantic_training_maps_room_file_list_)
            {
                std::cout << "   " << i << std::endl;
                cv::Mat training_map = cv::imread(package_path + i, 0);
                room_training_maps.push_back(training_map);
            }
            std::cout << "room_segmentation/semantic_training_maps_hallway_file_list = \n";
            std::vector <cv::Mat> hallway_training_maps;
            for (const auto & i : semantic_training_maps_hallway_file_list_)
            {
                std::cout << "   " << i << std::endl;
                cv::Mat training_map = cv::imread(package_path + i, 0);
                hallway_training_maps.push_back(training_map);
            }
            //train the algorithm
            semantic_segmentation.trainClassifiers(room_training_maps, hallway_training_maps,
                                                   classifier_path, load_semantic_features_);
        }

        AdaboostClassifier semantic_segmentation;
        semantic_segmentation.segmentMap(original_img, segmented_map, map_resolution, room_lower_limit_semantic_,
                                         room_upper_limit_semantic_,
                                         classifier_path, classifier_default_path,
                                         (display_segmented_map_ && DEBUG_DISPLAYS));
    }
    // set voronoi random field parameters
    if (room_segmentation_algorithm_ == 5 || train_vrf_)
    {
        room_upper_limit_voronoi_random_ = 10000.0;
        room_lower_limit_voronoi_random_ = 1.53;
        voronoi_random_field_epsilon_for_neighborhood_ = 7;
        min_neighborhood_size_ = 5;
        min_voronoi_random_field_node_distance_ = 7.0;
        max_voronoi_random_field_inference_iterations_ = 9000;
        max_iterations_ = 150;
        max_area_for_merging_ = 12.5;
        std::cout << "room_segmentation/room_area_factor_upper_limit = " << room_upper_limit_voronoi_random_
                  << std::endl;
        std::cout << "room_segmentation/room_area_factor_lower_limit = " << room_lower_limit_voronoi_random_
                  << std::endl;
        std::cout << "room_segmentation/voronoi_random_field_epsilon_for_neighborhood = "
                  << voronoi_random_field_epsilon_for_neighborhood_ << std::endl;
        std::cout << "room_segmentation/min_neighborhood_size = " << min_neighborhood_size_ << std::endl;
        std::cout << "room_segmentation/min_voronoi_random_field_node_distance = "
                  << min_voronoi_random_field_node_distance_ << std::endl;
        std::cout << "room_segmentation/max_voronoi_random_field_inference_iterations = "
                  << max_voronoi_random_field_inference_iterations_ << std::endl;
        std::cout << "room_segmentation/max_iterations = " << max_iterations_ << std::endl;
        std::cout << "room_segmentation/max_area_for_merging = " << max_area_for_merging_ << std::endl;
        // train the algorithm if wanted
        if (train_vrf_)
        {
            // voronoi random field segmentation method
            VoronoiRandomFieldSegmentation vrf_segmentation;
            // vector that stores the possible labels that are drawn in the training maps. Order: room - hallway - doorway
            std::vector <uint> possible_labels(3);
            possible_labels[0] = 77;
            possible_labels[1] = 115;
            possible_labels[2] = 179;
            // load the training maps
            std::vector <cv::Mat> training_maps;
            std::cout << "room_segmentation/vrf_training_maps_file_list_ = \n";
            for (const auto & i : vrf_training_maps_file_list_)
            {
                std::cout << "   " << i << std::endl << std::endl;
                cv::Mat training_map = cv::imread(package_path + i, 0);
                training_maps.push_back(training_map);
            }
            // load the voronoi maps
            std::vector <cv::Mat> voronoi_maps;
            std::cout << "room_segmentation/vrf_voronoi_maps_file_list_ = \n";
            for (const auto & i : vrf_voronoi_maps_file_list_)
            {
                std::cout << "   " << i << std::endl << std::endl;
                cv::Mat training_map = cv::imread(package_path + i, 0);
                voronoi_maps.push_back(training_map);
            }
            // load the voronoi-nodes maps
            std::vector <cv::Mat> voronoi_node_maps;
            std::cout << "room_segmentation/vrf_voronoi_node_maps_file_list_ = \n";
            for (const auto & i : vrf_voronoi_node_maps_file_list_)
            {
                std::cout << "   " << i << std::endl << std::endl;
                cv::Mat training_map = cv::imread(package_path + i, 0);
                voronoi_node_maps.push_back(training_map);
            }
            // load the original maps
            std::vector <cv::Mat> original_maps;
            std::cout << "room_segmentation/vrf_voronoi_node_maps_file_list = \n";
            for (const auto & i : vrf_original_maps_file_list_)
            {
                std::cout << "   " << i << std::endl << std::endl;
                cv::Mat training_map = cv::imread(package_path + i, 0);
                original_maps.push_back(training_map);
            }
            // train the algorithm
            vrf_segmentation.trainAlgorithms(original_maps, training_maps, voronoi_maps,
                                             voronoi_node_maps, possible_labels, classifier_storage_path,
                                             voronoi_random_field_epsilon_for_neighborhood_, max_iterations_,
                                             min_neighborhood_size_, min_voronoi_random_field_node_distance_);
        }
        VoronoiRandomFieldSegmentation vrf_segmentation;
        // vector that stores the possible labels that are drawn in the training maps. Order: room - hallway - doorway
        std::vector <uint> possible_labels(3);
        possible_labels[0] = 77;
        possible_labels[1] = 115;
        possible_labels[2] = 179;
        doorway_points_.clear();
        vrf_segmentation.segmentMap(original_img, segmented_map, voronoi_random_field_epsilon_for_neighborhood_,
                                    max_iterations_,
                                    min_neighborhood_size_, possible_labels, min_voronoi_random_field_node_distance_,
                                    (display_segmented_map_ && DEBUG_DISPLAYS), classifier_storage_path,
                                    classifier_default_path, max_voronoi_random_field_inference_iterations_,
                                    map_resolution, room_lower_limit_voronoi_random_, room_upper_limit_voronoi_random_,
                                    max_area_for_merging_, &doorway_points_);
    }
    // set passthrough parameters
    if (room_segmentation_algorithm_ == 99)
    {
        room_upper_limit_passthrough_ = 10000000.0;
        room_lower_limit_passthrough_ = 0.0;
        std::cout << "room_segmentation/room_area_factor_upper_limit = " << room_upper_limit_passthrough_ << std::endl;
        std::cout << "room_segmentation/room_area_factor_lower_limit = " << room_lower_limit_passthrough_ << std::endl;

        // pass through segmentation: takes a map which is already separated into unconnected areas and
        //                returns these as the resulting segmentation in the format of this program
        // todo: closing operation explicitly for bad maps --> needs parameterization
        // original_img.convertTo(segmented_map, CV_32SC1, 256, 0);		// occupied space = 0, free space = 65280
        cv::Mat original_img_eroded, temp;
        cv::erode(original_img, temp, cv::Mat(), cv::Point(-1, -1), 3);
        cv::dilate(temp, original_img_eroded, cv::Mat(), cv::Point(-1, -1), 3);
        original_img_eroded.convertTo(segmented_map, CV_32SC1, 256, 0);        // occupied space = 0, free space = 65280
        int label_index = 1;
//		cv::imshow("original_img", original_img_eroded);
//		cv::waitKey();
        for (int y = 0; y < segmented_map.rows; y++)
        {
            for (int x = 0; x < segmented_map.cols; x++)
            {
                // if original map is occupied space here or if the segmented map has already received a label for that cell --> skip
                if (original_img_eroded.at<uchar>(y, x) != 255 || segmented_map.at<int>(y, x) != 65280) continue;
                // fill each room area with a unique id
                cv::Rect rect;
                cv::floodFill(segmented_map, cv::Point(x, y), label_index, &rect, 0, 0, 4);
                // determine filled area
                double area = 0;
                for (int v = rect.y; v < segmented_map.rows; v++)
                    for (int u = rect.x; u < segmented_map.cols; u++)
                        if (segmented_map.at<int>(v, u) == label_index) area += 1.;
                area = map_resolution * map_resolution * area;    // convert from cells to m^2
                // exclude too small and too big rooms
                if (area < room_lower_limit_passthrough_ || area > room_upper_limit_passthrough_)
                {
                    for (int v = rect.y; v < segmented_map.rows; v++)
                        for (int u = rect.x; u < segmented_map.cols; u++)
                            if (segmented_map.at<int>(v, u) == label_index)
                                segmented_map.at<int>(v, u) = 0;
                }
                else label_index++;
            }
        }
        std::cout << "Labeled " << label_index - 1 << " segments." << std::endl;
    }
    //	looping_rate.sleep();
    // get the min/max-values and the room-centers
    // compute room label codebook
    // maps each room label to a position in the rooms vector
    std::map<int, size_t> label_vector_index_codebook;
    size_t vector_index = 0;
    for (int v = 0; v < segmented_map.rows; ++v)
    {
        for (int u = 0; u < segmented_map.cols; ++u)
        {
            const int label = segmented_map.at<int>(v, u);
            // do not count walls/obstacles or free space as label
            if (label > 0 && label < 65280)
            {
                if (label_vector_index_codebook.find(label) == label_vector_index_codebook.end())
                {
                    label_vector_index_codebook[label] = vector_index;
                    vector_index++;
                }
            }
        }
    }
    //min/max y/x-values vector for each room. Initialized with extreme values
    std::vector<int> min_x_value_of_the_room(label_vector_index_codebook.size(), 100000000);
    std::vector<int> max_x_value_of_the_room(label_vector_index_codebook.size(), 0);
    std::vector<int> min_y_value_of_the_room(label_vector_index_codebook.size(), 100000000);
    std::vector<int> max_y_value_of_the_room(label_vector_index_codebook.size(), 0);
    //vector of the central Point for each room, initially filled with Points out of the map
    std::vector<int> room_centers_x_values(label_vector_index_codebook.size(), -1);
    std::vector<int> room_centers_y_values(label_vector_index_codebook.size(), -1);
    //***********************Find min/max x and y coordinate and center of each found room********************
    // check y/x-value for every Pixel and make the larger/smaller value to the current value of the room
    for (int y = 0; y < segmented_map.rows; ++y)
    {
        for (int x = 0; x < segmented_map.cols; ++x)
        {
            const int label = segmented_map.at<int>(y, x);
            // if Pixel is white or black it is no room --> doesn't need to be checked
            if (label > 0 && label < 65280)
            {
                const int index = label_vector_index_codebook[label];
                min_x_value_of_the_room[index] = std::min(x, min_x_value_of_the_room[index]);
                max_x_value_of_the_room[index] = std::max(x, max_x_value_of_the_room[index]);
                max_y_value_of_the_room[index] = std::max(y, max_y_value_of_the_room[index]);
                min_y_value_of_the_room[index] = std::min(y, min_y_value_of_the_room[index]);
            }
        }
    }
    // get centers for each room
//	for (size_t idx = 0; idx < room_centers_x_values.size(); ++idx)
//	{
//		if (max_x_value_of_the_room[idx] != 0 && max_y_value_of_the_room[idx] != 0 && min_x_value_of_the_room[idx] != 100000000 && min_y_value_of_the_room[idx] != 100000000)
//		{
//			room_centers_x_values[idx] = (min_x_value_of_the_room[idx] + max_x_value_of_the_room[idx]) / 2;
//			room_centers_y_values[idx] = (min_y_value_of_the_room[idx] + max_y_value_of_the_room[idx]) / 2;
//			cv::circle(segmented_map, cv::Point(room_centers_x_values[idx], room_centers_y_values[idx]), 2, cv::Scalar(200*256), CV_FILLED);
//		}
//	}
    // use distance transform and mean shift to find good room centers that are reachable by the robot
    // first check whether a robot radius shall be applied to obstacles in order to exclude room center points that are not reachable by the robot
    cv::Mat segmented_map_copy = segmented_map;
    // stores for each pixel whether a path to another rooms exists for a robot of size robot_radius
    cv::Mat connection_to_other_rooms = cv::Mat::zeros(segmented_map.rows, segmented_map.cols, CV_8UC1);
    if (robot_radius > 0.0)
    {
        // consider robot radius for exclusion of non-reachable points
        segmented_map_copy = segmented_map.clone();
        cv::Mat map_8u, eroded_map;
        segmented_map_copy.convertTo(map_8u, CV_8UC1, 1., 0.);
        int number_of_erosions = (robot_radius / map_resolution);
        cv::erode(map_8u, eroded_map, cv::Mat(), cv::Point(-1, -1), number_of_erosions);
        for (int v = 0; v < segmented_map_copy.rows; ++v)
            for (int u = 0; u < segmented_map_copy.cols; ++u)
                if (eroded_map.at<uchar>(v, u) == 0)
                    segmented_map_copy.at<int>(v, u) = 0;
        // compute connectivity of remaining accessible room cells to other rooms
        bool stop = false;
        while (!stop)
        {
            stop = true;
            for (int v = 1; v < segmented_map_copy.rows - 1; ++v)
            {
                for (int u = 1; u < segmented_map_copy.cols - 1; ++u)
                {
                    // skip already identified cells
                    if (connection_to_other_rooms.at<uchar>(v, u) != 0) continue;
                    // only consider cells labeled as a room
                    const int label = segmented_map_copy.at<int>(v, u);
                    if (label <= 0 || label >= 65280) continue;
                    for (int dv = -1; dv <= 1; ++dv)
                    {
                        for (int du = -1; du <= 1; ++du)
                        {
                            if (dv == 0 && du == 0) continue;
                            const int neighbor_label = segmented_map_copy.at<int>(v + dv, u + du);
                            if (neighbor_label > 0 && neighbor_label < 65280 &&
                                (neighbor_label != label || (neighbor_label == label &&
                                                             connection_to_other_rooms.at<uchar>(v + dv, u + du) == 255)))
                            {
                                // either the room cell has a direct border to a different room or the room cell has a neighbor from the same room label with a connecting path to another room
                                connection_to_other_rooms.at<uchar>(v, u) = 255;
                                stop = false;
                            }
                        }
                    }
                }
            }
        }
    }
    // compute the room centers
    MeanShift2D ms;
    for (auto & it : label_vector_index_codebook)
    {
        // use robot_radius to avoid room centers that are not accessible by a robot with a given radius
        int trial = 1;
        if (robot_radius <= 0.) trial = 2;
        for (; trial <= 2; ++trial)
        {
            // compute distance transform for each room on the room cells that have some connection to another room (trial 1) or just on all cells of that room (trial 2)
            const int label = it.first;
            int number_room_pixels = 0;
            cv::Mat room = cv::Mat::zeros(segmented_map_copy.rows, segmented_map_copy.cols, CV_8UC1);
            for (int v = 0; v < segmented_map_copy.rows; ++v)
                for (int u = 0; u < segmented_map_copy.cols; ++u)
                    if (segmented_map_copy.at<int>(v, u) == label &&
                        (trial == 2 || connection_to_other_rooms.at<uchar>(v, u) == 255))
                    {
                        room.at<uchar>(v, u) = 255;
                        ++number_room_pixels;
                    }
            if (number_room_pixels == 0) continue;
            cv::Mat distance_map;
            // variable for the distance-transformed map, type: CV_32FC1
            cv::distanceTransform(room, distance_map, CV_DIST_L2, 5);
            // find point set with largest distance to obstacles
            double min_val = 0., max_val = 0.;
            cv::minMaxLoc(distance_map, &min_val, &max_val);
            std::vector <cv::Vec2d> room_cells;
            for (int v = 0; v < distance_map.rows; ++v)
                for (int u = 0; u < distance_map.cols; ++u)
                    if (distance_map.at<float>(v, u) > max_val * 0.95f)
                        room_cells.emplace_back(u, v);
            if (room_cells.empty())
                continue;
            // use meanshift to find the modes in that set
            cv::Vec2d room_center = ms.findRoomCenter(room, room_cells, map_resolution);
            const int index = it.second;
            room_centers_x_values[index] = room_center[0];
            room_centers_y_values[index] = room_center[1];
            if (!room_cells.empty()) break;
        }
    }
    // convert the segmented map into an indexed map which labels the segments with consecutive numbers (instead of arbitrary unordered labels in segmented map)
    cv::Mat indexed_map = segmented_map.clone();
    for (int y = 0; y < segmented_map.rows; ++y)
    {
        for (int x = 0; x < segmented_map.cols; ++x)
        {
            const int label = segmented_map.at<int>(y, x);
            //start value from 1 --> 0 is reserved for obstacles
            if (label > 0 && label < 65280)
                indexed_map.at<int>(y, x) = label_vector_index_codebook[label] + 1;
        }
    }
    if (display_segmented_map_)
    {
        // colorize the segmented map with the indices of the room_center vector
        cv::Mat color_segmented_map = indexed_map.clone();
        color_segmented_map.convertTo(color_segmented_map, CV_8U);
//        cv::cvtColor(color_segmented_map, color_segmented_map, CV_GRAY2BGR);
        cv::cvtColor(color_segmented_map, color_segmented_map, cv::COLOR_GRAY2BGR);
        for (size_t i = 1; i <= room_centers_x_values.size(); ++i)
        {
            //choose random color for each room
            const cv::Vec3b color((rand() % 250) + 1, (rand() % 250) + 1, (rand() % 250) + 1);
            for (size_t v = 0; v < indexed_map.rows; ++v)
                for (size_t u = 0; u < indexed_map.cols; ++u)
                    if (indexed_map.at<int>(v, u) == i)
                        color_segmented_map.at<cv::Vec3b>(v, u) = color;
        }
//		cv::Mat disp = segmented_map.clone();
        for (size_t index = 0; index < room_centers_x_values.size(); ++index)
#if CV_MAJOR_VERSION <= 3
                cv::circle(color_segmented_map, cv::Point(room_centers_x_values[index], room_centers_y_values[index]),
                           2, cv::Scalar(256), CV_FILLED);
#else
        cv::circle(color_segmented_map, cv::Point(room_centers_x_values[index], room_centers_y_values[index]), 2, cv::Scalar(256), cv::FILLED);
#endif
        cv::imshow("segmentation", color_segmented_map);
        cv::waitKey(0);
    }
    return 0;
}
